
#include "gen_task_node_utils.h"
#include "virtual_topo.h"


namespace hccl {

GenTaskNodeGraphBase::~GenTaskNodeGraphBase()
{
    for (auto& ele : toDeleteTaskNodeResource_) {
        if (ele == nullptr) {
            continue;
        }
        delete ele;
    }
    TaskQueueStub::Global()->Reset();
}

void GenTaskNodeGraphBase::LinkNode(TaskNodePtr parent, TaskNodePtr node)
{
    parent->children.push_back(node);
    node->parents.push_back(parent);
    return;
}

TaskNodePtr GenTaskNodeGraphBase::CreateHeadNode()
{
    head = new TaskNode(nullptr, -1, 0, 0);
    toDeleteTaskNodeResource_.push_back(head);
    return head;
}

TaskNodePtr GenTaskNodeGraphBase::CreatePostNode(RankId myRank, RankId remoteRankId, LinkInfo &link, u32 topicId,
                                                 u32 queueId, u32 pos, Stream& stream, NotifyTypeStub notifyType)
{
    TaskStub* post = new TaskStubPost(remoteRankId, link, topicId, notifyType);
    TaskNodePtr postNode = new TaskNode(post, myRank, queueId, pos);

    std::shared_ptr<TaskStub> spPost(post);
    TaskQueueStub::AppendTask(myRank, &stream, spPost);

    toDeleteTaskNodeResource_.push_back(postNode);
    return postNode;
}

TaskNodePtr GenTaskNodeGraphBase::CreateWaitNode(RankId myRank, RankId remoteRankId, LinkInfo &link, u32 topicId,
                                                 u32 queueId, u32 pos,  Stream& stream, NotifyTypeStub notifyType)
{
    TaskStub* wait = new TaskStubWait(remoteRankId, link, topicId, notifyType);
    TaskNodePtr waitNode = new TaskNode(wait, myRank, queueId, pos);

    std::shared_ptr<TaskStub> spWait(wait);
    TaskQueueStub::AppendTask(myRank, &stream, spWait);

    toDeleteTaskNodeResource_.push_back(waitNode);
    return waitNode;
}

TaskNodePtr GenTaskNodeGraphBase::CreateLocalPostToNode(u32 topicId, RankId myRank, u32 queueId, u32 pos, Stream& stream)
{
    TaskStub* postTo = new TaskStubLocalPostTo(topicId);
    TaskNodePtr postToNode = new TaskNode(postTo, myRank, queueId, pos);

    std::shared_ptr<TaskStub> spPostTo(postTo);
    TaskQueueStub::AppendTask(myRank, &stream, spPostTo);

    toDeleteTaskNodeResource_.push_back(postToNode);
    return postToNode;
}

TaskNodePtr GenTaskNodeGraphBase::CreateLocalWaitFromNode(u32 topicId, RankId myRank, u32 queueId, u32 pos, Stream& stream)
{
    TaskStub* waitFrom = new TaskStubLocalWaitFrom(topicId);
    TaskNodePtr waitFromNode = new TaskNode(waitFrom, myRank, queueId, pos);

    std::shared_ptr<TaskStub> spWaitFrom(waitFrom);
    TaskQueueStub::AppendTask(myRank, &stream, spWaitFrom);

    toDeleteTaskNodeResource_.push_back(waitFromNode);
    return waitFromNode;
}

TaskNodePtr GenTaskNodeGraphBase::CreateLocalCopyNode(DataSlice& srcSlice, DataSlice& dstSlice, RankId rankId,
                                                      u32 queueId, u32 pos, Stream& stream)
{
    TaskStub* localCopy = new TaskStubLocalCopy(srcSlice, dstSlice);
    TaskNodePtr localCopyNode = new TaskNode(localCopy, rankId, queueId, pos);

    std::shared_ptr<TaskStub> spLocalCopy(localCopy);
    TaskQueueStub::AppendTask(rankId, &stream, spLocalCopy);

    toDeleteTaskNodeResource_.push_back(localCopyNode);
    return localCopyNode;
}

TaskNodePtr GenTaskNodeGraphBase::CreateLocalReduceNode(DataSlice& srcSlice, DataSlice& dstSlice, RankId rankId,
                                                        u32 queueId, u32 pos, HcclDataType dataType, HcclReduceOp op, Stream& stream)
{
    TaskStub* localReduce = new TaskStubLocalReduce(srcSlice, dstSlice, dataType, op);
    TaskNodePtr localReduceNode = new TaskNode(localReduce, rankId, queueId, pos);

    std::shared_ptr<TaskStub> spLocalReduce(localReduce);
    TaskQueueStub::AppendTask(rankId, &stream, spLocalReduce);

    toDeleteTaskNodeResource_.push_back(localReduceNode);
    return localReduceNode;
}

TaskNodePtr GenTaskNodeGraphBase::CreateReadNode(DataSlice& localSlice, DataSlice& remoteSlice, RankId remoteRankId,
                                                 LinkInfo& link, RankId rankId, u32 queueId, u32 pos, Stream& stream)
{
    TaskStub* read = new TaskStubRead(remoteRankId, link, localSlice, remoteSlice);
    TaskNodePtr readNode = new TaskNode(read, rankId, queueId, pos);

    std::shared_ptr<TaskStub> spRead(read);
    TaskQueueStub::AppendTask(rankId, &stream, spRead);

    toDeleteTaskNodeResource_.push_back(readNode);
    return readNode;
}

TaskNodePtr GenTaskNodeGraphBase::CreateWriteNode(DataSlice& localSlice, DataSlice& remoteSlice, RankId remoteRankId,
                                                  LinkInfo& link, RankId rankId, u32 queueId, u32 pos, Stream& stream)
{
    TaskStub* write = new TaskStubWrite(remoteRankId, link, localSlice, remoteSlice);
    TaskNodePtr writeNode = new TaskNode(write, rankId, queueId, pos);

    std::shared_ptr<TaskStub> spWrite(write);
    TaskQueueStub::AppendTask(rankId, &stream, spWrite);

    toDeleteTaskNodeResource_.push_back(writeNode);
    return writeNode;
}

TaskNodePtr GenTaskNodeGraphBase::CreateReadReduceNode(DataSlice& localSlice, DataSlice& remoteSlice, RankId remoteRankId,
    LinkInfo& link, HcclDataType dataType, HcclReduceOp op, RankId rankId, u32 queueId, u32 pos, Stream& stream)
{
    TaskStub* readReduce = new TaskStubReadReduce(remoteRankId, link, localSlice, remoteSlice, dataType, op);
    TaskNodePtr readReduceNode = new TaskNode(readReduce, rankId, queueId, pos);

    std::shared_ptr<TaskStub> spReadReduce(readReduce);
    TaskQueueStub::AppendTask(rankId, &stream, spReadReduce);

    toDeleteTaskNodeResource_.push_back(readReduceNode);
    return readReduceNode;
}

TaskNodePtr GenTaskNodeGraphBase::CreateWriteReduceNode(DataSlice& localSlice, DataSlice& remoteSlice, RankId remoteRankId,
    LinkInfo& link, HcclDataType dataType, HcclReduceOp op, RankId rankId, u32 queueId, u32 pos, Stream& stream)
{
    TaskStub* writeReduce = new TaskStubWriteReduce(remoteRankId, link, localSlice, remoteSlice, dataType, op);
    TaskNodePtr writeReduceNode = new TaskNode(writeReduce, rankId, queueId, pos);

    std::shared_ptr<TaskStub> spWriteReduce(writeReduce);
    TaskQueueStub::AppendTask(rankId, &stream, spWriteReduce);

    toDeleteTaskNodeResource_.push_back(writeReduceNode);
    return writeReduceNode;
}

TaskNodePtr GenTaskNodeGraphBase::GenGraph()
{
    // 生成空的头部节点
    head = CreateHeadNode();
    return head;
}
}